using System;
using System.Collections.Generic;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

/*
 * 基础消息类，所有ComfyUI消息的基类
 * 任务执行完成时，消息发送顺序为execution_success-->status-->-->executing(data.node=null)
 */


namespace RSJWYFamework.Runtime
{
    /// <summary>
    /// ComfyUI websocket 消息转换器
    /// </summary>
    public class ComfyWebsocketMessageConverter : JsonConverter
    {
        // 消息类型字符串 → 对应的 Message 类型
        private static readonly Dictionary<string, Type> TypeMap = new()
        {
            [nameof(ComfyUIWebsocketMessageType.Status)]            = typeof(StatusMessage),
            [nameof(ComfyUIWebsocketMessageType.ExecutionStart)]   = typeof(ExecutionStartMessage),
            [nameof(ComfyUIWebsocketMessageType.Executing)]         = typeof(ExecutingMessage),
            [nameof(ComfyUIWebsocketMessageType.Progress)]          = typeof(ProgressMessage),
            [nameof(ComfyUIWebsocketMessageType.ExecutionSuccess)] = typeof(ExecutionSuccessMessage),
            [nameof(ComfyUIWebsocketMessageType.ExecutionCached)]  = typeof(ExecutionCachedMessage),
            // 以后再加
        };

        /// <summary>
        /// 判断是否可以转换为指定的消息类型
        /// </summary>
        /// <param name="objectType"></param>
        /// <returns></returns>
        public override bool CanConvert(Type objectType) => typeof(ComfyUIWebsocketBaseMessage<BaseComfyUImanager>).IsAssignableFrom(objectType);
        /// <summary>
        /// 从 JSON 读取消息
        /// </summary>
        /// <param name="reader"></param>
        /// <param name="objectType"></param>
        /// <param name="existingValue"></param>
        /// <param name="serializer"></param>
        /// <returns></returns>
        public override object ReadJson(JsonReader reader, Type objectType,
            object existingValue, JsonSerializer serializer)
        {
            var jo = JObject.Load(reader);

            // 取 type 字段
            var typeToken = jo["type"];
            if (typeToken == null)
                throw new JsonException("Missing 'type' field.");

            var typeStr = typeToken.Value<string>()!;
            if (!TypeMap.TryGetValue(typeStr, out var targetType))
                throw new JsonException($"Unknown message type: {typeStr}");

            // 让 JsonSerializer 继续反序列化到具体子类
            return jo.ToObject(targetType, serializer);
        }
        /// <summary>
        /// 写入 JSON 数据
        /// </summary>
        /// <param name="writer"></param>
        /// <param name="value"></param>
        /// <param name="serializer"></param>
        public override void WriteJson(JsonWriter writer, object value,
            JsonSerializer serializer)
        {
            // 如果用不到序列化，直接抛异常即可
            throw new NotImplementedException("Only deserialization is supported.");
        }
    }
    public enum ComfyUIWebsocketMessageType
    {
        /// <summary>
        /// 状态消息
        /// </summary>
        Status,
        /// <summary>
        /// 任务开始消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
        /// </summary>
        ExecutionStart,
        /// <summary>
        /// 任务执行中消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
        /// </summary>
        Executing,
        /// <summary>
        /// 任务步数数据，哪个节点的步数执行进度
        /// </summary>
        Progress,
        /// <summary>
        /// 任务成功消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
        /// </summary>
        ExecutionSuccess,
        /// <summary>
        /// 任务缓存消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
        /// </summary>
        ExecutionCached,
    }
    public abstract class ComfyUIWebsocketBaseMessage<TData> where TData : BaseComfyUImanager,new() // 约束TData必须有默认构造函数，方便反序列化
    {
        [JsonProperty("type")]
        public abstract string MsgType { get; } // 抽象属性，由子类指定具体type值

        [JsonProperty("data")]
        public TData MsgData { get; set; } = new TData(); // 泛型Data，默认初始化避免空引用
    }

    public class BaseComfyUImanager
    {
    }
    /// <summary>
    /// 状态消息
    /// </summary>
    public class StatusMessage : ComfyUIWebsocketBaseMessage<StatusData>
    {
        public override string MsgType => "status"; // 固定type值
    }
    /// <summary>
    /// 任务开始消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
    /// </summary>
    public class ExecutionStartMessage : ComfyUIWebsocketBaseMessage<ExecutionStartData>
    {
        public override string MsgType => "execution_start"; // 固定type值
    }
    
    /// <summary>
    /// 任务执行中消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
    /// </summary>
    public class ExecutingMessage : ComfyUIWebsocketBaseMessage<ExecutingData>
    {
        public override string MsgType => "executing"; // 固定type值
    }
   
    /// <summary>
    /// 任务步数数据，哪个节点的步数执行进度
    /// </summary>
    public class ProgressMessage : ComfyUIWebsocketBaseMessage<ProgressData>
    {
        public override string MsgType => "progress"; // 固定type值
    }
    
    /// <summary>
    /// 任务成功消息，携带返回你提交的任务ID，请判断任务ID是否与你提交的任务ID一致
    /// </summary>
    public class ExecutionSuccessMessage : ComfyUIWebsocketBaseMessage<ExecutionSuccessData>
    {
        public override string MsgType => "execution_success"; // 固定type值
    }
    /// <summary>
    /// 某些节点因为输出已存在于缓存中，本次执行被跳过。
    /// 缓存的节点列表，这些节点将在本次运行中跳过（上次执行后参数未变，直接使用上次缓存的结果）
    /// </summary>
    public class ExecutionCachedMessage : ComfyUIWebsocketBaseMessage<ExecutionCachedData>
    {
        public override string MsgType => "execution_cached"; // 固定type值
    }

    #region status
    public class StatusData : BaseComfyUImanager
    {
        [JsonProperty("status")]
        public StatusInfo Status { get; set; } = new StatusInfo();
        /// <summary>
        /// 会话ID
        /// </summary>
        /// <remarks>这是ComfyUI的会话ID，用于唯一标识一个会话，即你WS请求的ClientID</remarks>
        [JsonProperty("sid")]
        public string SID { get; set; } = string.Empty;
    }

    public class StatusInfo
    {
        [JsonProperty("exec_info")]
        public ExecInfo ExecInfo { get; set; } = new ExecInfo();
    }

    public class ExecInfo
    {
        /// <summary>
        /// 队列剩余数量
        /// </summary>
        /// <remarks>这是本ComfyUI的全局队列剩余数量，详细队列信息请请求GET /queue</remarks>
        [JsonProperty("queue_remaining")]
        public int QueueRemaining { get; set; }
    }


    #endregion
    #region execution_start
    
    public class ExecutionStartData : BaseComfyUImanager
    {
        /// <summary>
        /// 提交的图片的任务ID
        /// </summary>
        [JsonProperty("prompt_id")]
        public string PromptId { get; set; } = string.Empty;
        /// <summary>
        /// 消息时间戳，本消息发送时间的时间戳，单位毫秒
        /// </summary>
        [JsonProperty("prompt")]
        public long TimeStamp { get; set; }
    }

    #endregion

    #region executing
    public class ExecutingData : BaseComfyUImanager
    {
        /// <summary>
        /// 提交的图片的任务ID
        /// </summary>
        [JsonProperty("prompt_id")]
        public string PromptId { get; set; } = string.Empty;
        /// <summary>
        /// 当前节点执行到的节点，当为null时，同时代表任务执行完成
        /// </summary>
        [JsonProperty("node")]
        public string Node { get; set; }
        /// <summary>
        /// 当前节点执行到的节点的显示名称？
        /// 不知道，看发来信息是有这条数据，设置为可null，不需要管这个参数
        /// </summary>
        [JsonProperty("display_node")] 
        public string DisplayNode { get; set; }
    }

    #endregion
    #region execution_cached
    /// <summary>
    /// 任务缓存数据
    /// </summary>
    public class ExecutionCachedData : BaseComfyUImanager
    {
        /// <summary>
        /// 缓存的节点列表，这些节点将在本次运行中跳过（上次执行后参数未变，直接使用上次缓存的结果）
        /// </summary>
        [JsonProperty("nodes")]
        public string[] Nodes { get; set; }
        /// <summary>
        /// 提交的图片的任务ID
        /// </summary>
        [JsonProperty("prompt_id")]
        public string PromptId { get; set; } = string.Empty;
        /// <summary>
        /// 消息时间戳，本消息发送时间的时间戳，单位毫秒
        /// </summary>
        [JsonProperty("timestamp")]
        public long TimeStamp { get; set; }
    }
    
    #endregion
    

    #region progress
    public class ProgressData : BaseComfyUImanager
    {
        /// <summary>
        /// 当前节点执行的步数
        /// </summary>
        [JsonProperty("value")]
        public int Value { get; set; }
        
        /// <summary>
        /// 最大执行步数
        /// </summary>
        [JsonProperty("max")]
        public int Max { get; set; }
        
        /// <summary>
        /// 任务ID
        /// </summary>
        [JsonProperty("prompt_id")]
        public string PromptId { get; set; }
        
        
        /// <summary>
        /// 当前节点执行到的节点
        /// </summary>
        [JsonProperty("node")]
        public string Node { get; set; }
    }

    #endregion
    
    #region execution_success
    /// <summary>
    /// 任务成功数据
    /// </summary>
    public class ExecutionSuccessData : BaseComfyUImanager
    {
        /// <summary>
        /// 提交的图片的任务ID
        /// </summary>
        [JsonProperty("prompt_id")]
        public string PromptId { get; set; } = string.Empty;
        /// <summary>
        /// 消息时间戳，本消息发送时间的时间戳，单位毫秒
        /// </summary>
        [JsonProperty("prompt")]
        public long TimeStamp { get; set; }
    }
    #endregion
}